#!/usr/bin/env python

#
# S.M.A.C.K - An operating system kernel
# Copyright (C) 2010,2011 Mattias Holm and Kristian Rietveld
# For licensing and a full list of authors of the kernel, see the files
# COPYING and AUTHORS.
#


import unittest
import os
import hashlib		# for MD5

wfstool = "../wfstool"
test_image = "test.img"
test_files = ["test0.dat", "test1.dat", "test2.dat", "test3.dat", "test4.dat" ]

class TestCopy(unittest.TestCase):
	# Utilities
	#
	def put_file(self, local, remote):
		command = "%s %s put %s %s" % (wfstool, test_image, local, remote)
		self.assertEqual (os.system (command), 0)

	def get_file(self, remote, local):
		command = "%s %s get %s %s" % (wfstool, test_image, remote, local)
		self.assertEqual (os.system (command), 0)

	def remove_file(self, remote):
		command = "%s %s remove %s" % (wfstool, test_image, remote)
		self.assertEqual (os.system (command), 0)

	def add_directory(self, name):
		command = "%s %s adddir %s" % (wfstool, test_image, name)
		self.assertEqual (os.system (command), 0)

	def md5(self, file):
		fh = open (file, "rb")
		contents = fh.read ()
		fh.close ()

		m = hashlib.md5 ()
		m.update (contents)

		return m.hexdigest ()

	def compare_files(self, a, b):
		self.assertEqual (self.md5 (a), self.md5 (b))

	def check_files(self, files):
		for file in files:
			copy = file + ".copy"

			self.get_file (file, copy)
			self.compare_files (file, copy)
			os.unlink (copy)

	def fsck(self):
		command = "%s %s fsck 2> /dev/null" % (wfstool, test_image)
		status = os.system (command)

		if status != 0:
			print "fsck failed"

			# Run again with output enabled
			command = "%s %s fsck" % (wfstool, test_image)
			os.system (command)

			os.system ("mv test.img test.img.fail")

			self.fail ("fsck failed")

	# setUp & tearDown
	def setUp(self):
		command = "%s %s create" % (wfstool, test_image)
		self.assertEqual (os.system (command), 0)

		self.fsck ()

	def tearDown(self):
		os.unlink (test_image)

	#
	# Tests
	#

	# Put files
	#
	def test_put_single(self):
		file = test_files[0]
		copy = test_files[0] + ".copy"

		self.put_file (file, file)
		self.fsck ()

		self.get_file (file, copy)
		self.compare_files (file, copy)
		os.unlink (copy)

	def put_files(self, files):
		for file in files:
			self.put_file (file, file)

		self.fsck ()

		self.check_files (files)

		self.fsck()

	def test_put_two(self):
		files = [test_files[0], test_files[1]]
		self.put_files(files)

	def test_put_ascending(self):
		self.put_files(test_files)

	def test_put_descending(self):
		files = list(test_files)
		files.reverse()

		self.put_files(files)

	def test_put_shuffled(self):
		files = list()
		for i in [2, 4, 1, 3, 0]:
			files.append(test_files[i])

		self.put_files(files)

	# File enlargement on file overwrite
	#
	def file_overwrite(self, files):
		dest = "largefile"
		copy = dest + ".copy"

		for file in files:
			self.put_file (file, dest)
			self.fsck()

			self.get_file (dest, copy)
			self.compare_files (file, copy)
			os.unlink (copy)

		self.fsck()

	def test_file_overwrite_ascending(self):
		self.file_overwrite(test_files)

	def test_file_overwrite_descending(self):
		files = list(test_files)
		files.reverse()

		self.file_overwrite(files)

	def test_file_overwrite_shuffled(self):
		files = list()
		for i in [2, 4, 1, 3, 0]:
			files.append(test_files[i])

		self.file_overwrite(files)

	# File removal
	#
	def test_remove_ascending(self):
		self.test_put_ascending()

		check_files = list(test_files)
		for file in test_files:
			self.remove_file (file)
			self.fsck()

			check_files.remove(file)
			self.check_files(check_files)

		self.fsck()

	def test_remove_descending(self):
		self.test_put_ascending()

		files = list(test_files)
		files.reverse()

		check_files = list(files)
		for file in files:
			self.remove_file (file)
			self.fsck()

			check_files.remove(file)
			self.check_files(check_files)

		self.fsck()

	def test_remove_shuffle(self):
		self.test_put_ascending()

		files = list()
		for i in [2, 4, 1, 3, 0]:
			files.append(test_files[i])

		check_files = list(files)
		for file in files:
			self.remove_file (file)
			self.fsck()

			check_files.remove(file)
			self.check_files(check_files)

		self.fsck()

	# Directories
	#
	def test_add_directory(self):
		self.add_directory("test")

		self.fsck()

	def test_add_directories(self):
		self.add_directory("test")
		self.add_directory("test1")
		self.add_directory("test2")
		self.add_directory("test3")

		self.fsck()

	def test_remove_directory(self):
		self.add_directory("test")
		self.remove_file("test")

		self.fsck()

	def test_remove_directories(self):
		self.test_add_directories()

		self.remove_file("test2")
		self.remove_file("test1")
		self.remove_file("test")

		self.fsck()

	def test_directory_compound(self):
		self.test_remove_directories()

		self.add_directory("foo")
		self.add_directory("bar")

		self.remove_file("foo")

		self.fsck()

	# Compound tests
	#
	def test_compound0(self):
		self.test_put_ascending()
		self.test_remove_descending()
		self.test_file_overwrite_ascending()

	def test_compound1(self):
		self.test_put_ascending()
		self.test_put_ascending()
		self.test_remove_descending()
		self.test_file_overwrite_ascending()
		self.test_file_overwrite_ascending()

	def test_compound2(self):
		self.test_put_ascending()
		self.file_overwrite(test_files[0:3])
		self.test_remove_descending()
		self.file_overwrite(test_files[3:5])
		self.test_put_ascending()

	def test_compound3(self):
		self.put_files(test_files[0:3])
		self.file_overwrite(test_files[0:3])
		self.put_files(test_files[3:5])
		self.put_files(test_files[0:3])
		self.file_overwrite(test_files[3:5])
		self.put_files(test_files[0:3])
		self.test_remove_ascending()
		self.file_overwrite(test_files[2:4])

	def test_compound4(self):
		self.put_files(test_files[0:3])
		self.file_overwrite(test_files[0:3])
		self.put_files(test_files[3:5])
		self.test_add_directories()
		self.put_files(test_files[0:3])
		self.file_overwrite(test_files[3:5])
		self.remove_file("test2")
		self.remove_file("test")
		self.put_files(test_files[0:3])
		self.test_remove_ascending()
		self.file_overwrite(test_files[2:4])

if __name__ == '__main__':
	if not os.access (wfstool, os.R_OK | os.X_OK):
		print "error: compile wfstool first."
		exit (-1)

	for i in test_files:
		if not os.access (i, os.R_OK):
			print "error: file '%s' is missing." % (i)
			print "error: use gen-testfiles.sh to generate the testfiles"
			exit (-1)

	unittest.main ()
